package _12_avl_tree.case1;

import _00_my.annotation.Risky;

import java.util.Stack;

/**
 * @author: mornd
 * @dateTime: 2023/6/14 - 1:23
 * AVL 树
 * 二叉搜索树在插入或删除时，节点可能会失衡
 * 如果在插入或删除时通过旋转，始终让二叉搜索树保持平衡，那么便称之为自平衡的二叉搜索树
 * AVL 树是自平衡二叉搜索树的实现之一
 */

public class AVLTree<E> {
    public AVLNode root;

    /**
     * 添加
     *
     * @param k
     * @param value
     */
    public void put(int k, E value) {
        root = put(root, k, value);
    }

    private AVLNode put(AVLNode node, int k, E v) {
        if (node == null) {
            return new AVLNode(k, v);
        }
        if (k < node.key) {
            node.left = put(node.left, k, v);
        } else if (k > node.key) {
            node.right = put(node.right, k, v);
        } else {
            node.value = v;
            return node;
        }

        // 更新高度
        updateHeight(node);
        return balance(node); // 如果失衡，则重新平衡
    }

    /**
     * 移除节点
     * 1、node 为 null
     * 2、key 不存在
     * 3、找到，只有左孩子 或 只有右孩子
     * 4、找到，左右都有孩子，（后继节点是否相邻）
     * 5、更新高度
     * 6、自平衡
     *
     * 参考图片：img\删除前.png
     *
     * @param k
     */
    @Risky("参考remove2")
    public void remove(int k) {
        AVLNode curr = root;
        // 保存待删除节点的所有祖先
        Stack<AVLNode> stack = new Stack<>();

        while (curr != null) {
            if (k < curr.key) {
                stack.push(curr);
                curr = curr.left;
            } else if (k > curr.key) {
                stack.push(curr);
                curr = curr.right;
            } else {
                break;
            }
        }
        if(stack.isEmpty()) {
            root = null;
            return;
        }

        if (curr != null) { // != null 表示已找到带删除节点
            AVLNode parent = stack.peek();
            if (curr.left == null && curr.right == null) {
                shift(parent, curr, null);
                return;
            } else if (curr.left == null) {
                shift(parent, curr, curr.right);
            } else if (curr.right == null) {
                shift(parent, curr, curr.left);
            } else {
                // curr 的左右孩子都有

                // 找到后继节点
                AVLNode s = curr.right; // 后继节点
                AVLNode sParent = curr; // 后继节点的爹
                while (s.left != null) {
                    sParent = s;
                    s = s.left;
                }
                stack.push(s);

                if (sParent != curr) {  // 条件成立则不相邻

                    // 添加后继节点的所有父节点
                    s = curr.right;
                    sParent = curr;
                    while (s.left != null) {
                        sParent = s;
                        stack.push(sParent);
                        s = s.left;
                    }

                    shift(sParent, s, s.right);
                    stack.push(sParent);
                    s.right = curr.right;
                }

                shift(parent, curr, s);
                s.left = curr.left;
            }

            // 更新所有祖先的高度和保证自平衡
            while (!stack.isEmpty()) {
                AVLNode pop = stack.pop();
                updateHeight(pop);
                AVLNode balance = balance(pop);
                if(!stack.isEmpty()) {
                    shift(stack.peek(), pop, balance);
                }
            }
        }
    }

    // 托孤
    private void shift(AVLNode parent, AVLNode curr, AVLNode child) {
        if (parent == null) {
            root = child;
        } else if (parent.left == curr) {
            parent.left = child;
        } else if (parent.right == curr) {
            parent.right = child;
        }
    }

    // 删除方式二 递归实现
    public void remove2(int k) {
        root = remove2(root, k);
    }

    private AVLNode remove2(AVLNode node, int k) {
        if (node == null) {
            return null;
        }
        if (k < node.key) {
            node.left = remove2(node.left, k);
        } else if (k > node.key) {
            node.right = remove2(node.right, k);
        } else {
            if (node.left == null && node.right == null) {
                return null;
            } else if (node.left == null) {
                node = node.right;
            } else if (node.right == null) {
                node = node.left;
            } else {
                // 找后继节点
                AVLNode s = node.right;
                if (s.left != null) {
                    s = s.left;
                }
                // 递归处理后继节点
                s.right = remove2(node.right, s.key);
                s.left = node.left;
                node = s;
            }
        }

        // 更新高度
        updateHeight(node);
        // 自平衡
        return balance(node);
    }


    /**
     * 检查节点是否失衡，并重新平衡
     *
     * @param node
     * @return
     */
    private AVLNode balance(AVLNode node) {
        if (node == null) {
            return null;
        }
        int bf = balanceFactor(node);

        // 代码能进任一 if 块中表示已经失衡
        if (bf > 1 && balanceFactor(node.left) >= 0) { // LL
            // 左左失衡
            return rightRotate(node);
        } else if (bf > 1 && balanceFactor(node.left) < 0) { // LR
            // 左右失衡
            return leftRightRotate(node);
        } else if (bf < -1 && balanceFactor(node.right) > 0) { // RR
            return rightLeftRotate(node);
        } else if (bf < -1 && balanceFactor(node.right) <= 0) { // RL
            return leftRotate(node);
        }
        // 无需旋转
        return node;
    }

    /**
     * 左旋
     *
     * @param node 要旋转的节点
     * @return 新的根节点
     */
    private AVLNode leftRotate(AVLNode node) {
        AVLNode x = node.right;
        AVLNode t = x.left;

        x.left = node;
        node.right = t;

        // 更新高度
        updateHeight(node);
        updateHeight(x);

        return x;
    }

    private AVLNode rightRotate(AVLNode node) {
        AVLNode x = node.left;
        AVLNode t = x.right;

        x.right = node;
        node.left = t;

        updateHeight(node);
        updateHeight(x);

        return x;
    }

    /**
     * 左子节点先左旋，根节点再右旋
     *
     * @param node
     * @return
     */
    private AVLNode leftRightRotate(AVLNode node) {
        node.left = leftRotate(node.left);
        return rightRotate(node);
    }

    private AVLNode rightLeftRotate(AVLNode node) {
        node.right = rightRotate(node.right);
        return leftRotate(node);
    }

    /**
     * 获取平衡因子
     * 结果返回 1,0,-1 表示左右平衡
     * >1 || < -1 表示不平衡
     * >1 左边太高
     * <-1 右边太高
     *
     * @param node
     */
    private int balanceFactor(AVLNode node) {
        /*
            结果返回 1,0,-1 表示平衡
            >1 || < -1 表示不平衡

         */
        return height(node.left) - height(node.right);
    }

    /**
     * 更新节点高度
     */
    private void updateHeight(AVLNode node) {
        int l = height(node.left);
        int r = height(node.right);
        node.height = 1 + Math.max(l, r);
    }

    /**
     * 求节点的高度
     */
    public int height(AVLNode node) {
        return node == null ? 0 : node.height;
    }

    public class AVLNode {
        public int key;
        public E value;
        public AVLNode left;
        public AVLNode right;
        // 高度，默认1
        public int height = 1;

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, E value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key, int height) {
            this.key = key;
            this.height = height;
        }
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        traverse(root, "", "", sb);
        return sb.toString();
    }

    // 中序遍历打印树结构
    private void traverse(AVLNode node, String prefix, String childrenPrefix, StringBuilder sb) {
        if (node != null) {
            traverse(node.right, childrenPrefix + "├── ", childrenPrefix + "│   ", sb);
            sb.append(prefix).append(node.key).append("\n");
            traverse(node.left, childrenPrefix + "└── ", childrenPrefix + "    ", sb);
        }
    }
}
